import heapq
from itertools import count, islice
import math
import random
import networkx as nx
from networkx import Graph
import matplotlib as plt

class QKDNetwork:
    class LinkCapacity:
        def __init__(self, startNode, endNode, capacity):
            self.startNode = startNode
            self.endNode = endNode
            self.capacity = capacity

    class Node:
        def __init__(self, nodeID):
            self.nodeID = nodeID
            self.adjacentNodes = []

    def __init__(self, path = "", useNetworkx = False, networkxGraph: Graph = None):
        self.nodeList = []
        self.linkListCapacity: list[QKDNetwork.LinkCapacity] = []
        if(not useNetworkx):
            self.readDataFromTxt(path)
        else:
            if(networkxGraph == None):
                print("当你设置useNetworkx为True时，请指定networkxGraph字段")
            else:
                self.networkxGraph = networkxGraph
                self.readDataFromNetworkx(networkxGraph)
                self.initLinks()
                self.calculateCandidatePaths()
                self.calculateBrokenLinkFixPaths()
    
    def initLinks(self):
        self.numLinks = len(self.networkxGraph.edges)
        self.links = {}
        for i in range(self.numLinks):
            source, target = self.networkxGraph.edges[i]
            self.links[i] = (source, target)
    
    def calculateBrokenLinkFixPaths(self):
        self.brokenLinkFixPaths = {}
        for brokenLink in self.linkListCapacity:
            source, target = brokenLink.startNode, brokenLink.endNode
            paths = list(
                islice(nx.shortest_simple_paths(self.networkxGraph, source, target), 6)
            )
            # print((source, target), paths)
            self.brokenLinkFixPaths[(source, target)] = paths[1:]
                
    def calculateCandidatePaths(self):
        self.candidatePaths = {}
        self.maxCandidatePathLength = 0
        self.sd_num = 0
        self.sd = {}
        # 使用 Yen 算法为每个SD对计算5条候选路径
        for source in range(len(self.nodeList)):
            for target in range(source, len(self.nodeList)):
                if(source == target):
                    continue
                self.sd[self.sd_num] = (source, target)
                self.sd_num += 1
                paths = list(
                    islice(nx.shortest_simple_paths(self.networkxGraph, source, target), 5)
                )
                for path in paths:
                    pathLength = len(path)
                    if pathLength > self.maxCandidatePathLength:
                        self.maxCandidatePathLength = pathLength
                # print((source, target), paths)
                self.candidatePaths[(source, target)] = paths
        
    def getAllPathFromDFS(self, startNodeID, endNodeID, maxLength):
        def dfsHelper(currentNode, endNode, maxLength, currentPath, visited, paths):
            visited.add(currentNode)
            currentPath.append(currentNode)
            if currentNode == endNode:
                paths.append(list(currentPath))
            elif len(currentPath) <= maxLength:
                for adjacentNode in self.getNode(currentNode).adjacentNodes:
                    if adjacentNode not in visited:
                        dfsHelper(adjacentNode, endNode, maxLength, currentPath, visited, paths)
            currentPath.pop()
            visited.remove(currentNode)
        startNode = self.getNode(startNodeID).nodeID
        endNode = self.getNode(endNodeID).nodeID
        if startNode == -1 or endNode == -1:
            return [] 
        visited = set()
        paths = []
        dfsHelper(startNode, endNode, maxLength, [], visited, paths)
        return paths
        
    def getShortedPathWithOneBrokenLinkAndFlowMatrix(self, linkStartNodeID, linkEndNodeID, flowMatrix):
        distances = {}
        paths = {}
        visited = set()
        for node in self.nodeList:
            distances[node.nodeID] = float('-inf')  # 用负无穷初始化最大容量
            paths[node.nodeID] = []
        distances[linkStartNodeID] = float('inf')  # 起始节点的容量为正无穷
        while len(visited) < len(self.nodeList):
            max_capacity = float('-inf')
            max_node = None
            for node in self.nodeList:
                if node.nodeID not in visited and distances[node.nodeID] > max_capacity:
                    max_capacity = distances[node.nodeID]
                    max_node = node
            if max_node is None:
                break
            visited.add(max_node.nodeID)
            for neighborID in max_node.adjacentNodes:
                # 跳过破损的链路
                if (max_node.nodeID == linkStartNodeID and neighborID == linkEndNodeID) or (max_node.nodeID == linkEndNodeID and neighborID == linkStartNodeID):
                    continue
                neighbor = self.getNode(neighborID)
                link_capacity = self.getLinkCapacity(max_node.nodeID, neighbor.nodeID)
                flow = flowMatrix[max_node.nodeID-1][neighbor.nodeID-1]
                available_capacity = link_capacity - flow  # 允许负数
                new_capacity = min(distances[max_node.nodeID], available_capacity)
                if new_capacity > distances[neighbor.nodeID]:
                    distances[neighbor.nodeID] = new_capacity
                    paths[neighbor.nodeID] = paths[max_node.nodeID] + [neighbor.nodeID]
        return [linkStartNodeID, *paths[linkEndNodeID]] if paths[linkEndNodeID] else []
        
    def getShortedPathWithOneBrokenLink(self, linkStartNodeID, linkEndNodeID):
        # 创建一个字典来保存节点的距离和路径
        distances = {}
        paths = {}
        visited = set()
        # 初始化距离为无穷大，路径为空列表
        for node in self.nodeList:
            distances[node.nodeID] = float('inf')
            paths[node.nodeID] = []
        # 将起始节点的距离设置为0
        distances[linkStartNodeID] = 0

        while len(visited) < len(self.nodeList):
            # 找到未访问节点中距离最小的节点
            min_distance = float('inf')
            min_node = None
            for node in self.nodeList:
                if node.nodeID not in visited and distances[node.nodeID] < min_distance:
                    min_distance = distances[node.nodeID]
                    min_node = node
            if min_node is None:
                # 没有找到可到达的节点
                break
            # 标记节点为已访问
            visited.add(min_node.nodeID)
            # 更新相邻节点的距离和路径
            for neighborID in min_node.adjacentNodes:
                # 跳过破损的链路
                if (min_node.nodeID == linkStartNodeID and neighborID == linkEndNodeID) or (min_node.nodeID == linkEndNodeID and neighborID == linkStartNodeID):
                    continue
                neighbor = self.getNode(neighborID)
                # 计算新的距离
                new_distance = distances[min_node.nodeID] + 1
                if new_distance < distances[neighbor.nodeID]:
                    # 更新距离和路径
                    distances[neighbor.nodeID] = new_distance
                    paths[neighbor.nodeID] = paths[min_node.nodeID] + [neighbor.nodeID]

        return [linkStartNodeID, *paths[linkEndNodeID]] if distances[linkEndNodeID] != float('inf') else []

    # flowMatrix
    def getGreedyPathWithFlowMatrix(self, startNodeID, endNodeID, flowMatrix):
        # nodeNum = len(n.nodeList)
        # flowMatrix = [[0] * nodeNum for _ in range(nodeNum)]
        distances = {}
        paths = {}
        visited = set()
        for node in self.nodeList:
            distances[node.nodeID] = float('-inf')  # 用负无穷初始化最大容量
            paths[node.nodeID] = []
        distances[startNodeID] = float('inf')  # 起始节点的容量为正无穷
        while len(visited) < len(self.nodeList):
            max_capacity = float('-inf')
            max_node = None
            for node in self.nodeList:
                if node.nodeID not in visited and distances[node.nodeID] > max_capacity:
                    max_capacity = distances[node.nodeID]
                    max_node = node
            if max_node is None:
                break
            visited.add(max_node.nodeID)
            for neighborID in max_node.adjacentNodes:
                neighbor = self.getNode(neighborID)
                link_capacity = self.getLinkCapacity(max_node.nodeID, neighbor.nodeID)
                flow = flowMatrix[max_node.nodeID-1][neighbor.nodeID-1]
                available_capacity = link_capacity - flow  # 允许负数
                new_capacity = min(distances[max_node.nodeID], available_capacity)
                if new_capacity > distances[neighbor.nodeID]:
                    distances[neighbor.nodeID] = new_capacity
                    paths[neighbor.nodeID] = paths[max_node.nodeID] + [neighbor.nodeID]
        return [startNodeID, *paths[endNodeID]] if paths[endNodeID] else []

    def getShortedPath(self, startNodeID, endNodeID):
        # 创建一个字典来保存节点的距离和路径
        distances = {}
        paths = {}
        visited = set()
        # 初始化距离为无穷大，路径为空列表
        for node in self.nodeList:
            distances[node.nodeID] = float('inf')
            paths[node.nodeID] = []
        # 将起始节点的距离设置为0
        distances[startNodeID] = 0
        while len(visited) < len(self.nodeList):
            # 找到未访问节点中距离最小的节点
            min_distance = float('inf')
            min_node = None
            for node in self.nodeList:
                if node.nodeID not in visited and distances[node.nodeID] < min_distance:
                    min_distance = distances[node.nodeID]
                    min_node = node
            if min_node is None:
                # 没有找到可到达的节点
                break
            # 标记节点为已访问
            visited.add(min_node.nodeID)
            # 更新相邻节点的距离和路径
            for neighborID in min_node.adjacentNodes:
                neighbor = self.getNode(neighborID)
                # 计算新的距离
                new_distance = distances[min_node.nodeID] + 1
                if new_distance < distances[neighbor.nodeID]:
                    # 更新距离和路径
                    distances[neighbor.nodeID] = new_distance
                    paths[neighbor.nodeID] = paths[min_node.nodeID] + [neighbor.nodeID]
        return paths[endNodeID]
        
    def showNetwork(self):
        with open('generate.txt', 'r') as file:
            lines = file.readlines()
        n = int(lines[0])
        nodes = [int(line) for line in lines[1:n+1]]
        G = nx.Graph()
        G.add_nodes_from(nodes)
        m = int(lines[n+1])
        for line in lines[n+2:]:
            start, end, weight = map(int, line.split())
            G.add_edge(start, end, weight=weight)
        fig, ax = plt.subplots(figsize=(10, 10))
        pos = nx.spring_layout(G, seed=41) 
        nx.draw(G, pos, ax=ax, with_labels=True, node_color='red', node_size=300, font_size=8, alpha=1)
        edge_labels = nx.get_edge_attributes(G, 'weight')
        nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=6)
        plt.title("Network Topology")
        plt.show()
        # plt.savefig("topo.svg", format="svg")
    
    def getLinkNums(self):
        count = 0
        for node in self.nodeList:
            for _ in node.adjacentNodes:
                count += 1
        return count
    
    def getNeighborNodeIDs(self, nodeID):
        node = self.getNode(nodeID)
        return node.adjacentNodes
    
    def isIJNLeverageL(self, i, j, n, k, l):
        sd = self.sd[i]
        path = self.candidatePaths[sd][j]
        start = path[n]
        end = path[n + 1]
        link_l = self.links[l]
        fixPath = self.brokenLinkFixPaths[(start, end)][k]
        for i in range(len(fixPath)-1):
            if fixPath[i] == link_l[0] and fixPath[i+1] == link_l[1]:
                return True
            if fixPath[i] == link_l[1] and fixPath[i+1] == link_l[0]:
                return True
        return False
    
    def isIJLeverageL(self, i, j, l):
        # 请求i的第j条candidate path是否利用了l作为中继路径
        sd = self.sd[i]
        path = self.candidatePaths[sd][j]
        link = self.links[l]
        for i in range(len(path)-1):
            if path[i] == link[0] and path[i+1] == link[1]:
                return True
            if path[i] == link[1] and path[i+1] == link[0]:
                return True
        return False

    def getLinkCapacityByLinkID(self, linkID):
        edge = self.networkxGraph.edges[linkID]
        return edge['weight']
    
    def getLinkCapacity(self, startID, endID):
        for linkCapacity in self.linkListCapacity:
            if linkCapacity.startNode == startID and linkCapacity.endNode == endID:
                return linkCapacity.capacity
            if linkCapacity.startNode == endID and linkCapacity.endNode == startID:
                return linkCapacity.capacity
        return -1
    
    def setLinkCapacity(self, startID, endID, capacity):
        self.linkListCapacity.append(self.LinkCapacity(startID, endID, capacity))

    def getNode(self, nodeId):
        if nodeId - 1 >= len(self.nodeList) or nodeId - 1 < 0:
            return self.Node(-1)
        return self.nodeList[nodeId - 1]

    def readDataFromNetworkx(self, graph: Graph):
        for nodeID in range(1, graph.number_of_nodes()+1):
            self.nodeList.append(self.Node(nodeID))
        for (u, v) in graph.edges():
            self.nodeList[u - 1].adjacentNodes.append(v)
            self.nodeList[v - 1].adjacentNodes.append(u)
            self.linkListCapacity.append(self.LinkCapacity(u, v, graph.get_edge_data(u, v)['weight']))

    def readDataFromTxt(self, path):
        with open(path, 'r') as file:
            lines = file.readlines()
            numNodes = int(lines[0])
            nodeIDs = [int(line.strip()) for line in lines[1 : numNodes + 1]]
            
            for nodeID in nodeIDs:
                self.nodeList.append(self.Node(nodeID))
            
            numLinks = int(lines[numNodes + 1])
            for i in range(numNodes + 2, numNodes + numLinks + 2):
                linkData = lines[i].split()
                startNode = int(linkData[0])
                endNode = int(linkData[1])
                self.nodeList[startNode - 1].adjacentNodes.append(endNode)
                self.nodeList[endNode - 1].adjacentNodes.append(startNode)
                capacity = int(linkData[2])
                linkCapacity = self.LinkCapacity(startNode, endNode, capacity)
                self.linkListCapacity.append(linkCapacity)
